import math
from random import sample, shuffle

import cv2
import numpy as np
from PIL import Image
from tensorflow import keras

from utils.utils import cvtColor, preprocess_input


class YoloDatasets(keras.utils.Sequence):
    def __init__(self, annotation_lines, input_shape, anchors, batch_size, num_classes, anchors_mask, epoch_now,
                 epoch_length, mosaic, train, mosaic_ratio=0.7):
        self.annotation_lines = annotation_lines
        self.length = len(self.annotation_lines)

        self.input_shape = input_shape
        self.anchors = anchors
        self.batch_size = batch_size
        self.num_classes = num_classes
        self.anchors_mask = anchors_mask
        self.epoch_now = epoch_now - 1
        self.epoch_length = epoch_length
        self.mosaic = mosaic
        self.train = train
        self.mosaic_ratio = mosaic_ratio

        self.threshold = 4

    def __len__(self):
        return math.ceil(len(self.annotation_lines) / float(self.batch_size))

    def __getitem__(self, index):
        image_data = []
        box_data = []
        for i in range(index * self.batch_size, (index + 1) * self.batch_size):
            i = i % self.length
            # ---------------------------------------------------#
            #   训练时进行数据的随机增强
            #   验证时不进行数据的随机增强
            # ---------------------------------------------------#
            if self.mosaic:
                if self.rand() < 0.5 and self.epoch_now < self.epoch_length * self.mosaic_ratio:
                    lines = sample(self.annotation_lines, 3)
                    lines.append(self.annotation_lines[i])
                    shuffle(lines)
                    image, box = self.get_random_data_with_Mosaic(lines, self.input_shape)
                else:
                    image, box = self.get_random_data(self.annotation_lines[i], self.input_shape, random=self.train)
            else:
                image, box = self.get_random_data(self.annotation_lines[i], self.input_shape, random=self.train)
            image_data.append(preprocess_input(np.array(image, np.float32)))
            box_data.append(box)

        image_data = np.array(image_data)
        box_data = np.array(box_data)
        y_true = self.preprocess_true_boxes(box_data, self.input_shape, self.anchors, self.num_classes)
        return [image_data, *y_true], np.zeros(self.batch_size)

    def generate(self):
        i = 0
        while True:
            image_data = []
            box_data = []
            for b in range(self.batch_size):
                if i == 0:
                    np.random.shuffle(self.annotation_lines)

                print(self.annotation_lines[i])  ##测试用，用完要注释
                # ---------------------------------------------------#
                #   训练时进行数据的随机增强
                #   验证时不进行数据的随机增强
                # ---------------------------------------------------#
                if self.mosaic:
                    if self.rand() < 0.5 and self.epoch_now < self.epoch_length * self.mosaic_ratio:
                        lines = sample(self.annotation_lines, 3)
                        lines.append(self.annotation_lines[i])
                        shuffle(lines)
                        image, box = self.get_random_data_with_Mosaic(lines, self.input_shape)
                    else:
                        image, box = self.get_random_data(self.annotation_lines[i], self.input_shape, random=self.train)
                else:
                    image, box = self.get_random_data(self.annotation_lines[i], self.input_shape, random=self.train)

                i = (i + 1) % self.length
                image_data.append(preprocess_input(np.array(image, np.float32)))
                box_data.append(box)
            image_data = np.array(image_data)
            box_data = np.array(box_data)
            y_true = self.preprocess_true_boxes(box_data, self.input_shape, self.anchors, self.num_classes)
            yield image_data, y_true[0], y_true[1], y_true[2]

    def on_epoch_end(self):
        self.epoch_now += 1
        shuffle(self.annotation_lines)

    def rand(self, a=0, b=1):
        return np.random.rand() * (b - a) + a

    ##2022-11-16 改max_boxes： 500->105
    def get_random_data(self, annotation_line, input_shape, max_boxes=201, jitter=.3, hue=.1, sat=0.7, val=0.4,
                        random=True):
        line = annotation_line.split()
        # ------------------------------#
        #   读取图像并转换成RGB图像
        # ------------------------------#
        image = Image.open(line[0])
        image = cvtColor(image)
        # ------------------------------#
        #   获得图像的高宽与目标高宽
        # ------------------------------#
        iw, ih = image.size
        h, w = input_shape
        # ------------------------------#
        #   获得预测框
        # ------------------------------#
        box = np.array([np.array(list(map(int, box.split(',')))) for box in line[1:]])

        if not random:
            scale = min(w / iw, h / ih)
            nw = int(iw * scale)
            nh = int(ih * scale)
            dx = (w - nw) // 2
            dy = (h - nh) // 2

            # ---------------------------------#
            #   将图像多余的部分加上灰条
            # ---------------------------------#
            image = image.resize((nw, nh), Image.BICUBIC)
            new_image = Image.new('RGB', (w, h), (128, 128, 128))
            new_image.paste(image, (dx, dy))
            image_data = np.array(new_image, np.float32)

            # ---------------------------------#
            #   对真实框进行调整
            # ---------------------------------#
            box_data = np.zeros((max_boxes, 5))
            if len(box) > 0:
                np.random.shuffle(box)
                box[:, [0, 2]] = box[:, [0, 2]] * nw / iw + dx
                box[:, [1, 3]] = box[:, [1, 3]] * nh / ih + dy
                box[:, 0:2][box[:, 0:2] < 0] = 0
                box[:, 2][box[:, 2] >= w] = w - 1  # 减1 2023-02-10
                box[:, 3][box[:, 3] >= h] = h - 1  # 减1 2023-02-10
                box_w = box[:, 2] - box[:, 0]
                box_h = box[:, 3] - box[:, 1]
                box = box[np.logical_and(box_w > 1, box_h > 1)]
                if len(box) > max_boxes: box = box[:max_boxes]
                box_data[:len(box)] = box

            return image_data, box_data

        # ------------------------------------------#
        #   对图像进行缩放并且进行长和宽的扭曲
        # ------------------------------------------#
        new_ar = iw / ih * self.rand(1 - jitter, 1 + jitter) / self.rand(1 - jitter, 1 + jitter)
        # scale = self.rand(.25, 2)    ## self.rand(.25, 2)  这里要改为 (.5， 1.5) 不需要太大，
        scale = self.rand(.5, 1.5)  ## self.rand(.25, 2)  这里要改为 (.5， 1.5) 不需要太大，
        if new_ar < 1:
            nh = int(scale * h)
            nw = int(nh * new_ar)
        else:
            nw = int(scale * w)
            nh = int(nw / new_ar)
        image = image.resize((nw, nh), Image.BICUBIC)

        # ------------------------------------------#
        #   将图像多余的部分加上灰条
        # ------------------------------------------#
        dx = int(self.rand(0, w - nw))
        dy = int(self.rand(0, h - nh))
        new_image = Image.new('RGB', (w, h), (128, 128, 128))
        new_image.paste(image, (dx, dy))
        image = new_image

        # ------------------------------------------#
        #   翻转图像
        # ------------------------------------------#
        flip = self.rand() < .5
        if flip: image = image.transpose(Image.FLIP_LEFT_RIGHT)

        # nRotate = np.random.randint(0, 4)##增加旋转2023-02-10
        # if nRotate == 1:
        #     image = image.transpose(Image.ROTATE_90)
        # elif nRotate == 2:
        #     image = image.transpose(Image.ROTATE_180)
        # elif nRotate == 3:
        #     image = image.transpose(Image.ROTATE_270)

        image_data = np.array(image, np.uint8)
        # ---------------------------------#
        #   对图像进行色域变换
        #   计算色域变换的参数
        # ---------------------------------#
        r = np.random.uniform(-1, 1, 3) * [hue, sat, val] + 1
        # ---------------------------------#
        #   将图像转到HSV上
        # ---------------------------------#
        hue, sat, val = cv2.split(cv2.cvtColor(image_data, cv2.COLOR_RGB2HSV))
        dtype = image_data.dtype
        # ---------------------------------#
        #   应用变换
        # ---------------------------------#
        x = np.arange(0, 256, dtype=r.dtype)
        lut_hue = ((x * r[0]) % 180).astype(dtype)
        lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
        lut_val = np.clip(x * r[2], 0, 255).astype(dtype)

        image_data = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val)))
        image_data = cv2.cvtColor(image_data, cv2.COLOR_HSV2RGB)

        # ---------------------------------#
        #   对真实框进行调整
        # ---------------------------------#
        box_data = np.zeros((max_boxes, 5))
        # box_data = np.zeros((1, 5))##2022-11-16 觉得没必要
        if len(box) > 0:
            np.random.shuffle(box)
            box[:, [0, 2]] = box[:, [0, 2]] * nw / iw + dx
            box[:, [1, 3]] = box[:, [1, 3]] * nh / ih + dy

            isChange, ins, ins2 = self.isChangeLabel(box, nw, nh, w, h)  ##2022-11-16
            if isChange == True:
                box[ins, 4] = 2  ##修改标签为partial
                box = box[ins2]  ##剔除<0.1的box

            if flip: box[:, [0, 2]] = w - box[:, [2, 0]] - 1  ## 减1 2023-02-10
            box[:, 0:2][box[:, 0:2] < 0] = 0
            box[:, 2][box[:, 2] >= w] = w - 1  ## 减1 2023-02-10
            box[:, 3][box[:, 3] >= h] = h - 1  ## 减1 2023-02-10
            box_w = box[:, 2] - box[:, 0]
            box_h = box[:, 3] - box[:, 1]
            box = box[np.logical_and(box_w > 1, box_h > 1)]  # discard invalid box

            if len(box) > max_boxes: box = box[:max_boxes]
            box_data[:len(box)] = box
            # box_data = box##觉得上一步没必要，直接等于算了,原来是为了后面m个样本能够平成一个矩阵

        return image_data, box_data

    def merge_bboxes(self, bboxes, cutx, cuty):
        merge_bbox = []
        for i in range(len(bboxes)):
            for box in bboxes[i]:
                tmp_box = []
                x1, y1, x2, y2 = box[0], box[1], box[2], box[3]

                if i == 0:
                    if y1 > cuty or x1 > cutx:
                        continue
                    if y2 >= cuty and y1 <= cuty:
                        y2 = cuty
                    if x2 >= cutx and x1 <= cutx:
                        x2 = cutx

                if i == 1:
                    if y2 < cuty or x1 > cutx:
                        continue
                    if y2 >= cuty and y1 <= cuty:
                        y1 = cuty
                    if x2 >= cutx and x1 <= cutx:
                        x2 = cutx

                if i == 2:
                    if y2 < cuty or x2 < cutx:
                        continue
                    if y2 >= cuty and y1 <= cuty:
                        y1 = cuty
                    if x2 >= cutx and x1 <= cutx:
                        x1 = cutx

                if i == 3:
                    if y1 > cuty or x2 < cutx:
                        continue
                    if y2 >= cuty and y1 <= cuty:
                        y2 = cuty
                    if x2 >= cutx and x1 <= cutx:
                        x1 = cutx
                tmp_box.append(x1)
                tmp_box.append(y1)
                tmp_box.append(x2)
                tmp_box.append(y2)
                tmp_box.append(box[-1])
                merge_bbox.append(tmp_box)
        return merge_bbox

    def get_random_data_with_Mosaic(self, annotation_line, input_shape, max_boxes=105, jitter=0.3, hue=.1, sat=0.7,
                                    val=0.4):
        h, w = input_shape
        min_offset_x = self.rand(0.3, 0.7)
        min_offset_y = self.rand(0.3, 0.7)

        image_datas = []
        box_datas = []
        index = 0
        for line in annotation_line:
            # ---------------------------------#
            #   每一行进行分割
            # ---------------------------------#
            line_content = line.split()
            # ---------------------------------#
            #   打开图片
            # ---------------------------------#
            image = Image.open(line_content[0])
            image = cvtColor(image)

            # ---------------------------------#
            #   图片的大小
            # ---------------------------------#
            iw, ih = image.size
            # ---------------------------------#
            #   保存框的位置
            # ---------------------------------#
            box = np.array([np.array(list(map(int, box.split(',')))) for box in line_content[1:]])

            # ---------------------------------#
            #   是否翻转图片
            # ---------------------------------#
            flip = self.rand() < .5
            if flip and len(box) > 0:
                image = image.transpose(Image.FLIP_LEFT_RIGHT)
                box[:, [0, 2]] = iw - box[:, [2, 0]]

            # ------------------------------------------#
            #   对图像进行缩放并且进行长和宽的扭曲
            # ------------------------------------------#
            new_ar = iw / ih * self.rand(1 - jitter, 1 + jitter) / self.rand(1 - jitter, 1 + jitter)
            scale = self.rand(.4, 1)
            if new_ar < 1:
                nh = int(scale * h)
                nw = int(nh * new_ar)
            else:
                nw = int(scale * w)
                nh = int(nw / new_ar)
            image = image.resize((nw, nh), Image.BICUBIC)

            # -----------------------------------------------#
            #   将图片进行放置，分别对应四张分割图片的位置
            # -----------------------------------------------#
            if index == 0:
                dx = int(w * min_offset_x) - nw
                dy = int(h * min_offset_y) - nh
            elif index == 1:
                dx = int(w * min_offset_x) - nw
                dy = int(h * min_offset_y)
            elif index == 2:
                dx = int(w * min_offset_x)
                dy = int(h * min_offset_y)
            elif index == 3:
                dx = int(w * min_offset_x)
                dy = int(h * min_offset_y) - nh

            new_image = Image.new('RGB', (w, h), (128, 128, 128))
            new_image.paste(image, (dx, dy))
            image_data = np.array(new_image)

            index = index + 1
            box_data = []
            # ---------------------------------#
            #   对box进行重新处理
            # ---------------------------------#
            if len(box) > 0:
                np.random.shuffle(box)
                box[:, [0, 2]] = box[:, [0, 2]] * nw / iw + dx
                box[:, [1, 3]] = box[:, [1, 3]] * nh / ih + dy
                box[:, 0:2][box[:, 0:2] < 0] = 0
                box[:, 2][box[:, 2] > w] = w
                box[:, 3][box[:, 3] > h] = h
                box_w = box[:, 2] - box[:, 0]
                box_h = box[:, 3] - box[:, 1]
                box = box[np.logical_and(box_w > 1, box_h > 1)]
                box_data = np.zeros((len(box), 5))
                box_data[:len(box)] = box

            image_datas.append(image_data)
            box_datas.append(box_data)

        # ---------------------------------#
        #   将图片分割，放在一起
        # ---------------------------------#
        cutx = int(w * min_offset_x)
        cuty = int(h * min_offset_y)

        new_image = np.zeros([h, w, 3])
        new_image[:cuty, :cutx, :] = image_datas[0][:cuty, :cutx, :]
        new_image[cuty:, :cutx, :] = image_datas[1][cuty:, :cutx, :]
        new_image[cuty:, cutx:, :] = image_datas[2][cuty:, cutx:, :]
        new_image[:cuty, cutx:, :] = image_datas[3][:cuty, cutx:, :]

        new_image = np.array(new_image, np.uint8)
        # ---------------------------------#
        #   对图像进行色域变换
        #   计算色域变换的参数
        # ---------------------------------#
        r = np.random.uniform(-1, 1, 3) * [hue, sat, val] + 1
        # ---------------------------------#
        #   将图像转到HSV上
        # ---------------------------------#
        hue, sat, val = cv2.split(cv2.cvtColor(new_image, cv2.COLOR_RGB2HSV))
        dtype = new_image.dtype
        # ---------------------------------#
        #   应用变换
        # ---------------------------------#
        x = np.arange(0, 256, dtype=r.dtype)
        lut_hue = ((x * r[0]) % 180).astype(dtype)
        lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
        lut_val = np.clip(x * r[2], 0, 255).astype(dtype)

        new_image = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val)))
        new_image = cv2.cvtColor(new_image, cv2.COLOR_HSV2RGB)

        # ---------------------------------#
        #   对框进行进一步的处理
        # ---------------------------------#
        new_boxes = self.merge_bboxes(box_datas, cutx, cuty)

        # ---------------------------------#
        #   将box进行调整
        # ---------------------------------#
        box_data = np.zeros((max_boxes, 5))
        if len(new_boxes) > 0:
            if len(new_boxes) > max_boxes: new_boxes = new_boxes[:max_boxes]
            box_data[:len(new_boxes)] = new_boxes
        return new_image, box_data

    def get_near_points(self, x, y, i, j):
        sub_x = x - i
        sub_y = y - j
        if sub_x > 0.5 and sub_y > 0.5:
            return [[0, 0], [1, 0], [0, 1]]
        elif sub_x < 0.5 and sub_y > 0.5:
            return [[0, 0], [-1, 0], [0, 1]]
        elif sub_x < 0.5 and sub_y < 0.5:
            return [[0, 0], [-1, 0], [0, -1]]
        else:
            return [[0, 0], [1, 0], [0, -1]]

    def preprocess_true_boxes(self, true_boxes, input_shape, anchors, num_classes):
        assert (true_boxes[..., 4] < num_classes).all(), 'class id must be less than num_classes'
        # -----------------------------------------------------------#
        #   获得框的坐标和图片的大小
        # -----------------------------------------------------------#
        true_boxes = np.array(true_boxes, dtype='float32')
        input_shape = np.array(input_shape, dtype='int32')

        # -----------------------------------------------------------#
        #   一共有三个特征层数
        # -----------------------------------------------------------#
        num_layers = len(self.anchors_mask)
        # -----------------------------------------------------------#
        #   m为图片数量，grid_shapes为网格的shape
        # -----------------------------------------------------------#
        m = true_boxes.shape[0]
        grid_shapes = [input_shape // {0: 32, 1: 16, 2: 8}[l] for l in range(num_layers)]
        # -----------------------------------------------------------#
        #   y_true的格式为(m,13,13,3,85)(m,26,26,3,85)(m,52,52,3,85)
        # -----------------------------------------------------------#
        y_true = [np.zeros((m, grid_shapes[l][0], grid_shapes[l][1], len(self.anchors_mask[l]), 5 + num_classes),
                           dtype='float32') for l in range(num_layers)]
        # -----------------------------------------------------#
        #   anchors_best_ratio
        # -----------------------------------------------------#
        box_best_ratios = [np.zeros((m, grid_shapes[l][0], grid_shapes[l][1], len(self.anchors_mask[l])),
                                    dtype='float32') for l in range(num_layers)]

        # -----------------------------------------------------------#
        #   通过计算获得真实框的中心和宽高
        #   中心点(m,n,2) 宽高(m,n,2)
        # -----------------------------------------------------------#
        boxes_xy = (true_boxes[..., 0:2] + true_boxes[..., 2:4]) // 2
        boxes_wh = true_boxes[..., 2:4] - true_boxes[..., 0:2]
        # -----------------------------------------------------------#
        #   将真实框归一化到小数形式
        # -----------------------------------------------------------#
        true_boxes[..., 0:2] = boxes_xy / input_shape[::-1]
        true_boxes[..., 2:4] = boxes_wh / input_shape[::-1]

        # -----------------------------------------------------------#
        #   [9,2] -> [9,2]
        # -----------------------------------------------------------#
        anchors = np.array(anchors, np.float32)

        # -----------------------------------------------------------#
        #   长宽要大于0才有效
        # -----------------------------------------------------------#
        valid_mask = boxes_wh[..., 0] > 0

        for b in range(m):
            # -----------------------------------------------------------#
            #   对每一张图进行处理
            # -----------------------------------------------------------#
            wh = boxes_wh[b, valid_mask[b]]

            if len(wh) == 0: continue
            # -------------------------------------------------------#
            #   wh                      : num_true_box, 2
            #   anchors                 : 9, 2
            #
            #   ratios_of_gt_anchors    : num_true_box, 9, 2
            #   ratios_of_anchors_gt    : num_true_box, 9, 2
            #
            #   ratios                  : num_true_box, 9, 4
            #   max_ratios              : num_true_box, 9
            # -------------------------------------------------------#
            ratios_of_gt_anchors = np.expand_dims(wh, 1) / np.expand_dims(anchors, 0)
            ratios_of_anchors_gt = np.expand_dims(anchors, 0) / np.expand_dims(wh, 1)
            ratios = np.concatenate([ratios_of_gt_anchors, ratios_of_anchors_gt], axis=-1)
            max_ratios = np.max(ratios, axis=-1)

            for t, ratio in enumerate(max_ratios):
                # -------------------------------------------------------#
                #   ratio : 9
                # -------------------------------------------------------#
                over_threshold = ratio < self.threshold
                over_threshold[np.argmin(ratio)] = True
                # -----------------------------------------------------------#
                #   找到每个真实框所属的特征层
                # -----------------------------------------------------------#
                for l in range(num_layers):
                    for k, n in enumerate(self.anchors_mask[l]):
                        if not over_threshold[n]:
                            continue
                        # -----------------------------------------------------------#
                        #   floor用于向下取整，找到真实框所属的特征层对应的x、y轴坐标
                        # -----------------------------------------------------------#
                        i = np.floor(true_boxes[b, t, 0] * grid_shapes[l][1]).astype('int32')
                        j = np.floor(true_boxes[b, t, 1] * grid_shapes[l][0]).astype('int32')
                        offsets = self.get_near_points(true_boxes[b, t, 0] * grid_shapes[l][1],
                                                       true_boxes[b, t, 1] * grid_shapes[l][0], i, j)
                        for offset in offsets:
                            local_i = i + offset[0]
                            local_j = j + offset[1]

                            if local_i >= grid_shapes[l][1] or local_i < 0 or local_j >= grid_shapes[l][
                                0] or local_j < 0:
                                continue

                            if box_best_ratios[l][b, local_j, local_i, k] != 0:
                                if box_best_ratios[l][b, local_j, local_i, k] > ratio[n]:
                                    y_true[l][b, local_j, local_i, k, :] = 0
                                else:
                                    continue
                            # -----------------------------------------------------------#
                            #   c指的是当前这个真实框的种类
                            # -----------------------------------------------------------#
                            c = true_boxes[b, t, 4].astype('int32')
                            # -----------------------------------------------------------#
                            #   y_true的shape为(m,13,13,3,85)(m,26,26,3,85)(m,52,52,3,85)
                            #   最后的85可以拆分成4+1+80，4代表的是框的中心与宽高、
                            #   1代表的是置信度、80代表的是种类
                            # -----------------------------------------------------------#
                            y_true[l][b, local_j, local_i, k, 0:4] = true_boxes[b, t, 0:4]
                            y_true[l][b, local_j, local_i, k, 4] = 1
                            y_true[l][b, local_j, local_i, k, 5 + c] = 1
                            box_best_ratios[l][b, local_j, local_i, k] = ratio[n]

        return y_true

    def isChangeLabel(self, boxes, nw, nh, w, h, ratio_up=0.75, ratio_low=0.1):  # 2022-11-16
        ##不用管是否超出边界，后面有步骤去除，因为 矩形边 会出现负数被筛掉
        if (np.array([nw, nh]) <= np.array([w, h])).all():
            return False, None, None

        # ins1 = boxes[:, 0:2] >= 0  ##大部分图片都会超，这个判断就有点多余
        # if ins1.all():
        #     ins2 = boxes[:, 2:4] <= [w, h]
        #     if ins2.all():
        #         return False, None

        ww = [0, 0, w, h]
        wcoord = np.maximum(boxes[:, 0:2], ww[0:2])
        hcoord = np.minimum(boxes[:, 2:4], ww[2:4])

        inarea = np.prod(hcoord - wcoord, axis=1)
        area = np.prod(boxes[:, 2:4] - boxes[:, 0:2], axis=1)
        ratio = inarea / area

        ins2 = ratio >= ratio_low
        ins = np.logical_and(ratio < ratio_up, ins2)
        return True, ins, ins2

    ##纯碎测试用
    def testGenerate(self):
        i = 0
        while True:
            image_data = []
            box_data = []
            for b in range(self.batch_size):
                if i == 0:
                    np.random.shuffle(self.annotation_lines)
                # ---------------------------------------------------#
                #   训练时进行数据的随机增强
                #   验证时不进行数据的随机增强
                # ---------------------------------------------------#
                if self.mosaic:
                    if self.rand() < 0.5 and self.epoch_now < self.epoch_length * self.mosaic_ratio:
                        lines = sample(self.annotation_lines, 3)
                        lines.append(self.annotation_lines[i])
                        shuffle(lines)
                        image, box = self.get_random_data_with_Mosaic(lines, self.input_shape)
                    else:
                        image, box = self.get_random_data(self.annotation_lines[i], self.input_shape, random=self.train)
                else:
                    image, box = self.get_random_data(self.annotation_lines[i], self.input_shape, random=self.train)

                i = (i + 1) % self.length
                image_data.append(preprocess_input(np.array(image, np.float32)))
                box_data.append(box)

                # break##方便调试，只取一个image_data与 box_data

            image_data = np.array(image_data)
            box_data = np.array(box_data)
            y_true = self.preprocess_true_boxes(box_data, self.input_shape, self.anchors, self.num_classes)
            return image_data, y_true[0], y_true[1], y_true[2], box_data


## 2023-12-10 增加的训练集构成方式，增加每隔一段epoch就进行全体训练数据集的训练
class YoloDatasetNew(YoloDatasets):
    def __init__(self, annotation_lines, input_shape, anchors, batch_size, num_classes, anchors_mask, epoch_now,
                 epoch_length, mosaic, train, mosaic_ratio=0.7, args=None):

        super().__init__(annotation_lines, input_shape, anchors, batch_size, num_classes, anchors_mask, epoch_now,
                 epoch_length, mosaic, train, mosaic_ratio)  # 显式调用父类的初始化方法
        ##########
        if args is not None:
            self.data_change = True
            self.epoch_current = 0
            self.epoch_start = args["epoch_start"]
            self.epoch_duration = args["epoch_duration"]
            self.new_lines = args["new_lines"]
            self.true_lines = self.annotation_lines
            self.period = self.epoch_start + self.epoch_duration
        else:
            self.data_change = False
        #############

    def on_epoch_end(self):
        self.epoch_now += 1

        # 2023-12-10 额外处理
        if self.data_change:
            self.epoch_current += 1
            if self.epoch_current % self.period == self.epoch_start:
                self.annotation_lines = self.new_lines
            elif self.epoch_current % self.period == 1:
                self.annotation_lines = self.true_lines
        #####

        shuffle(self.annotation_lines)

